import matplotlib.pyplot as plt
import seaborn as sns
plt.rcParams["font.family"] = ["SimHei", "Microsoft YaHei", "SimSun", "sans-serif"]
plt.rcParams["axes.unicode_minus"] = False
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler, LabelEncoder
from sklearn.neighbors import KNeighborsClassifier
from sklearn.svm import SVC
# 新增模型导入
from sklearn.linear_model import LogisticRegression
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import RandomForestClassifier
# 模型评估和保存工具
from sklearn.metrics import classification_report, accuracy_score, confusion_matrix
from joblib import dump
import matplotlib.pyplot as plt
import seaborn as sns

class StockClassification:
    """
    股票分类模型：基于财务指标和价格波动，将股票分为四类（高收益高风险/低收益高风险/高收益低风险/低收益低风险）
    支持模型：KNN、SVM、逻辑回归、决策树、随机森林
    """
    
    def __init__(self, data_path='fina_indicator.csv'):
        """
        初始化：加载数据
        :param data_path: 数据文件路径（CSV格式）
        """
        self.data_path = data_path
        self.df = None  # 原始数据
        self.processed_df = None  # 处理后的数据
        self.features = None  # 特征矩阵
        self.target = None  # 目标变量
        self.scaler = StandardScaler()  # 标准化器
        self.label_encoder = LabelEncoder()  # 标签编码器
        self.models = {}  # 存储训练好的模型
        # 新增：定义要训练的所有模型（便于扩展和管理）
        self.model_configs = {
            'knn_model.joblib': KNeighborsClassifier(n_neighbors=5),
            'svm': SVC(kernel='rbf', probability=True, random_state=42),
            'logistic_reg': LogisticRegression(max_iter=1000, random_state=42),  # 逻辑回归
            'decision_tree': DecisionTreeClassifier(max_depth=8, random_state=42),  # 决策树
            'random_forest': RandomForestClassifier(n_estimators=100, max_depth=10, random_state=42)  # 随机森林
        }

    def load_data(self):
        """加载CSV数据"""
        try:
            self.df = pd.read_csv(self.data_path)
            print(f"成功加载数据，数据形状：{self.df.shape}")
            return self.df
        except Exception as e:
            print(f"数据加载失败：{e}")
            raise

    def preprocess_data(self):
        """数据预处理：计算比例指标、生成分类标签、特征选择"""
        if self.df is None:
            raise ValueError("请先调用load_data()加载数据")
        
        self.processed_df = self.df.copy()
        
        # 1. 计算收益潜力和风险潜力比例（基于价格数据）
        required_price_cols = ['max_close', 'min_close', 'the_close']
        if not all(col in self.processed_df.columns for col in required_price_cols):
            raise KeyError(f"数据缺少必要的价格列：{required_price_cols}")
        
        self.processed_df['max_ratio'] = self.processed_df['max_close'] / self.processed_df['the_close']  # 收益潜力（越高越好）
        self.processed_df['min_ratio'] = self.processed_df['min_close'] / self.processed_df['the_close']  # 风险潜力（越低风险越高）

        # 2. 自动计算分类阈值（基于分位数）
        high_return_threshold = self.processed_df['max_ratio'].quantile(0.4)  # 前40%为高收益
        high_risk_threshold = self.processed_df['min_ratio'].quantile(0.4)  # 前40%为高风险

        print(f"分类阈值 - 高收益阈值：{high_return_threshold:.3f}，高风险阈值：{high_risk_threshold:.3f}")

        # 3. 生成四类分类标签
        conditions = [
            (self.processed_df['max_ratio'] >= high_return_threshold) & (self.processed_df['min_ratio'] <= high_risk_threshold),
            (self.processed_df['max_ratio'] < high_return_threshold) & (self.processed_df['min_ratio'] <= high_risk_threshold),
            (self.processed_df['max_ratio'] >= high_return_threshold) & (self.processed_df['min_ratio'] > high_risk_threshold),
            (self.processed_df['max_ratio'] < high_return_threshold) & (self.processed_df['min_ratio'] > high_risk_threshold)
        ]
        choices = ['高收益高风险', '低收益高风险', '高收益低风险', '低收益低风险']
        self.processed_df['category'] = np.select(conditions, choices, default='未知')

        # 4. 特征选择（财务指标）
        feature_cols = [
            'eps', 'total_revenue_ps', 'undist_profit_ps', 'gross_margin',
            'fcff', 'fcfe', 'tangible_asset', 'bps', 'grossprofit_margin', 'npta'
        ]
        self.features = self.processed_df[feature_cols].copy()
        
        # 5. 处理缺失值（新增：避免模型训练报错）
        self.features = self.features.fillna(self.features.mean())  # 用均值填充缺失值
        
        # 6. 目标变量编码
        self.target = self.label_encoder.fit_transform(self.processed_df['category'])
        
        print(f"特征矩阵形状：{self.features.shape}")
        print(f"目标变量分布：{pd.Series(self.label_encoder.inverse_transform(self.target)).value_counts().to_dict()}")
        
        return self.features, self.target

    def train_models(self, test_size=0.2, random_state=42):
        """训练所有配置的模型（包括新增的逻辑回归、决策树、随机森林）"""
        if self.features is None or self.target is None:
            raise ValueError("请先调用preprocess_data()预处理数据")
        
        # 划分训练集和测试集（分层抽样，保证类别分布一致）
        X_train, X_test, y_train, y_test = train_test_split(
            self.features, self.target, test_size=test_size, random_state=random_state, stratify=self.target
        )
        
        # 特征标准化（对逻辑回归、SVM、KNN重要，对树模型影响小但不冲突）
        X_train_scaled = self.scaler.fit_transform(X_train)
        X_test_scaled = self.scaler.transform(X_test)

        # 训练所有模型（循环遍历配置的模型）
        print("\n=== 模型训练开始 ===")
        for model_name, model in self.model_configs.items():
            try:
                # 训练模型
                model.fit(X_train_scaled, y_train)
                self.models[model_name] = model  # 保存训练好的模型
                print(f"\n{model_name.upper()} 模型训练完成")
                
                # 模型评估
                y_pred = model.predict(X_test_scaled)
                accuracy = accuracy_score(y_test, y_pred)
                
                # 打印评估结果
                print(f"准确率：{accuracy:.4f}")
                print("分类报告：")
                print(classification_report(
                    y_test, y_pred, target_names=self.label_encoder.classes_, zero_division=0
                ))
                
                # 新增：绘制混淆矩阵（可视化模型表现）
                self.plot_confusion_matrix(y_test, y_pred, model_name)
                
            except Exception as e:
                print(f"{model_name.upper()} 模型训练失败：{e}")

        return self.models

    def plot_confusion_matrix(self, y_true, y_pred, model_name):
        """绘制混淆矩阵（可视化模型分类效果）"""
        plt.figure(figsize=(8, 6))
        cm = confusion_matrix(y_true, y_pred)
        sns.heatmap(
            cm, 
            annot=True, 
            fmt='d', 
            cmap='Blues',
            xticklabels=self.label_encoder.classes_,
            yticklabels=self.label_encoder.classes_
        )
        plt.title(f'{model_name.upper()} 混淆矩阵')
        plt.xlabel('预测标签')
        plt.ylabel('真实标签')
        plt.tight_layout()
        plt.show()

    def save_models(self, save_dir='./saved_models/'):
        """保存所有训练好的模型、标准化器和标签编码器"""
        import os
        os.makedirs(save_dir, exist_ok=True)
        
        # 保存每个模型
        for model_name, model in self.models.items():
            dump(model, os.path.join(save_dir, f'{model_name}_model.joblib'))
        
        # 保存标准化器和标签编码器
        dump(self.scaler, os.path.join(save_dir, 'scaler.joblib'))
        dump(self.label_encoder, os.path.join(save_dir, 'label_encoder.joblib'))
        
        print(f"\n所有模型已保存至：{save_dir}")

    def run_pipeline(self):
        """完整流程执行：加载数据 -> 预处理 -> 训练模型 -> 保存模型"""
        self.load_data()
        self.preprocess_data()
        self.train_models()
        self.save_models()
        print("\n=== 完整流程执行完成 ===")

    def predict(self, new_data):
        """使用训练好的模型进行预测（支持单条/多条数据）"""
        if not self.models:
            raise ValueError("请先调用train_models()训练模型")
        
        # 数据预处理（和训练时保持一致）
        new_data_processed = pd.DataFrame(new_data) if isinstance(new_data, dict) else new_data
        new_data_processed = new_data_processed.fillna(self.features.mean())  # 用训练集均值填充缺失值
        new_data_scaled = self.scaler.transform(new_data_processed)
        
        # 所有模型预测结果
        predictions = {}
        for model_name, model in self.models.items():
            pred_labels = model.predict(new_data_scaled)
            pred_labels_text = self.label_encoder.inverse_transform(pred_labels)  # 转为文本标签
            predictions[model_name] = pred_labels_text
        
        return predictions


# 数据查询辅助函数（保持原有功能）
def query_stock_data(cursor, ts_code, ann_date, new_list, row):
    """
    从数据库查询股票价格数据，计算价格指标并添加到结果列表
    :param cursor: 数据库游标
    :param ts_code: 股票代码
    :param ann_date: 公告日期
    :param new_list: 存储结果的列表
    :param row: 包含财务指标的行数据
    """
    query = f"""
        SELECT closes FROM stock_price 
        WHERE ts_code = '{ts_code}' 
        AND trade_date >= '{ann_date}'
        LIMIT 30  # 取公告后30天价格数据
    """
    
    try:
        cursor.execute(query)
        result = cursor.fetchall()
        
        df_price = pd.DataFrame(result, columns=['closes'])
        if len(df_price) > 0:
            max_close = df_price['closes'].max()
            min_close = df_price['closes'].min()
            the_close = df_price['closes'].iloc[-1]
            
            new_list.append({
                'ts_code': ts_code,
                'ann_date': ann_date,
                'max_close': max_close,
                'min_close': min_close,
                'the_close': the_close,
                'eps': row.get('eps', np.nan),
                'total_revenue_ps': row.get('total_revenue_ps', np.nan),
                'undist_profit_ps': row.get('undist_profit_ps', np.nan),
                'gross_margin': row.get('gross_margin', np.nan),
                'fcff': row.get('fcff', np.nan),
                'fcfe': row.get('fcfe', np.nan),
                'tangible_asset': row.get('tangible_asset', np.nan),
                'bps': row.get('bps', np.nan),
                'grossprofit_margin': row.get('grossprofit_margin', np.nan),
                'npta': row.get('npta', np.nan),
            })
    except Exception as e:
        print(f"查询股票{ts_code}数据失败：{e}")


# 主程序入口
if __name__ == "__main__":
    # 1. 执行完整分类流程（包含5个模型训练）
    classifier = StockClassification(data_path='fina_indicator.csv')
    classifier.run_pipeline()

    # 2. （可选）使用训练好的模型进行预测示例
    # 构造一条新的财务数据（特征顺序需和训练时一致）
    new_stock_data = pd.DataFrame([{
        'eps': 0.8,
        'total_revenue_ps': 5.2,
        'undist_profit_ps': 3.1,
        'gross_margin': 0.45,
        'fcff': 2.3,
        'fcfe': 1.8,
        'tangible_asset': 10.5,
        'bps': 4.2,
        'grossprofit_margin': 0.52,
        'npta': 0.08
    }])
    
    # 预测
    predictions = classifier.predict(new_stock_data)
    print("\n=== 新股票分类预测结果 ===")
    for model_name, pred in predictions.items():
        print(f"{model_name.upper()} 预测类别：{pred[0]}")